""" Graph Average Exploration - Figure 4"""

import os
import matplotlib
matplotlib.use('agg')
import matplotlib.style
matplotlib.style.use('seaborn-white')
matplotlib.rc('text', usetex=False)
matplotlib.rc('font', family='CMU Serif')
matplotlib.rc('font', serif='CMU Serif')
matplotlib.rc('font', size=13)
import matplotlib.pyplot as plt

DIR_DATA = 'data/'
DIR_RESULTS = 'results/'

YLIM_EXPLORATION     = 15
YTICKS_EXPLORATION   = [0, 3, 6, 9, 12, 15]
XTICKS_EXPLORATION   = [0, 100, 200, 300, 400, 500]
PLOT_FONT_SIZE       = 15
PLOT_AXIS_FONT_SIZE  = 15
PLOT_LINE_WIDTH      = 1.25
PLOT_AXIS_LABEL_PAD  = 10

def save_lst(lst, fname):
    fname = str(fname.strip())
    if lst is None: lst = list()
    delete_file(fname)
    with open(fname, 'w') as f:
        for i in range(len(lst)):
            item = str(lst[i])
            if i < (len(lst) - 1): item += "\n"
            f.write(item)

def load_lst(fname, force2int=False, must_exist=True):
    out = list()
    fname = str(fname.strip())
    if not os.path.exists(fname):
        if must_exist:
            raise IOError
        else:
            return out
    with open(fname, 'Ur', encoding='latin1') as f:  # universal newlines support; read-only
        lines = f.readlines()
        for line in lines:
            value = str(line).strip()
            if force2int: value = int(value)
            out.append(value)
    return out

def delete_file(fname):
    fname = str(fname.strip())
    try: os.remove(fname)
    except: pass


def plot_comparison(label_a, data_a, label_b, data_b, fname):
    """ Graph comparison of the average exploration for two history """
    plt.close('all')
    x = list(range(len(data_a)))
    fig, ax = plt.subplots()
    fig.set_size_inches(6.5, 5)
    ax.set_ylim(0, 15)
    plt.yticks(YTICKS_EXPLORATION)
    plt.xticks(XTICKS_EXPLORATION)
    if YLIM_EXPLORATION: ax.set_ylim(0,YLIM_EXPLORATION)
    if YTICKS_EXPLORATION: ax.set_yticks(YTICKS_EXPLORATION)
    ax.plot(x, data_a, lw=PLOT_LINE_WIDTH, ls='--', label=label_a, color='black',  alpha=.85)
    ax.plot(x, data_b, lw=PLOT_LINE_WIDTH, ls='-',  label=label_b, color='black', alpha=.85)
    ax.xaxis.grid(True, linestyle='-', color='grey', alpha=.25)
    ax.yaxis.grid(True, linestyle='-', color='grey', alpha=.25)
    ax.set_xlabel('Move Number',labelpad=PLOT_AXIS_LABEL_PAD, fontsize=PLOT_AXIS_FONT_SIZE)
    ax.set_ylabel('Number of Exploratory Actions',labelpad=PLOT_AXIS_LABEL_PAD, fontsize=PLOT_AXIS_FONT_SIZE)
    plt.legend(frameon=1, loc=2)
    frame = ax.legend_.get_frame()
    frame.set_facecolor('white')
    frame.set_edgecolor('white')
    if os.path.exists(fname): os.remove(fname)
    plt.savefig(fname, format='pdf')
    plt.close()


# COMMANDLINE  ===============================================================================================================================================
if __name__ == "__main__":

    HUMAN_RESULTS_1  = [float(i) for i in load_lst(DIR_DATA + 'avg-exploration-config-1.txt')]
    HUMAN_RESULTS_3  = [float(i) for i in load_lst(DIR_DATA + 'avg-exploration-config-3.txt')]
    HUMAN_RESULTS_10 = [float(i) for i in load_lst(DIR_DATA + 'avg-exploration-config-10.txt')]
    HUMAN_RESULTS_12 = [float(i) for i in load_lst(DIR_DATA + 'avg-exploration-config-12.txt')]

    plot_comparison(label_a='Losses', data_a=HUMAN_RESULTS_3, label_b='Gains', data_b=HUMAN_RESULTS_1, fname=DIR_RESULTS + '01_Graph_Figure4a.pdf')
    plot_comparison(label_a='Losses', data_a=HUMAN_RESULTS_12, label_b='Gains', data_b=HUMAN_RESULTS_10, fname=DIR_RESULTS + '01_Graph_Figure4b.pdf')

